"""
知识管理系统 - 分类搜索测试工具
"""

from db_utils import Database, VectorStore, Category


# 创建数据库连接和向量存储实例
db = Database()
vector_store = VectorStore(db)


def main():
    print("知识管理系统 - 分类搜索测试工具")
    print("============================")

    # 显示所有可用分类
    print("\n可用分类列表:")
    session = db.connect()
    categories = session.query(Category).all()
    for category in categories:
        print(f"  {category.id}. {category.name} - {category.description}")

    # 获取用户输入
    try:
        category_id = input("\n请输入要搜索的分类ID (留空表示搜索所有文档): ").strip()
        category_id = int(category_id) if category_id else None

        query_text = input("请输入搜索关键词: ")
        top_k = int(input("请输入要返回的结果数量 (默认5): ") or "5")

        # 执行搜索
        print("\n搜索中...")
        results = vector_store.search_similar_documents(
            query_text=query_text, top_k=top_k, category_id=category_id
        )

        # 显示结果
        category_str = f"分类{category_id}" if category_id else "所有分类"
        print(f"\n在{category_str}中搜索 '{query_text}' 的结果 ({len(results)} 个):")
        for i, (doc_id, content, title, similarity, category_ids) in enumerate(results):
            # 获取分类名称
            categories = (
                session.query(Category).filter(Category.id.in_(category_ids)).all()
            )
            category_names = [cat.name for cat in categories]

            print(f"\n{i+1}. 文档ID: {doc_id}")
            print(f"   标题: {title}")
            print(f"   相似度: {similarity:.4f}")
            print(
                f"   分类: {', '.join(category_names)} (ID: {', '.join(map(str, category_ids))})"
            )
            print(
                f"   内容摘要: {content[:100]}..."
                if len(content) > 100
                else f"   内容: {content}"
            )
    except ValueError:
        print("错误: 请输入有效的数字类型")
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        db.disconnect(session)


if __name__ == "__main__":
    main()
